import argparse
from typing import List
from basic_models.models.DGNS import DGN_R1, DGN_R2
from basic_models.models.Gacpn import GACPN_R1, GACPN_R2, GACPN_R3
from basic_models.models.commNet import CommNet

from basic_models.models.TranDGN import TranDGN, TranDGNWithoutMulti, TranDGNWithoutProcessor, TranDGNWithoutTran
from basic_models.envs.surviving import Surviving
from basic_models.envs.predator_prey_env import PredatorPreyEnv
from .buffer import ReplayBuffer


class ConfigSet:
    def __init__(self):
        self.__parser = argparse.ArgumentParser(prog="ConfigSet")
        self.__parser.add_argument('--hidden_dim', '-hd', type=int, default=128)
        self.__parser.add_argument('--i_episode', '-i_e', type=int, default=0)
        self.__parser.add_argument('--capacity', '-cap', type=int, default=20000)
        self.__parser.add_argument('--batch_size', '-b', type=int, default=256, dest='batch_size')
        self.__parser.add_argument('--n_epoch', '-ne', type=int, default=5, help='在多少个epoch后，进行目标网络的更新')
        self.__parser.add_argument('--epsilon', '-e', type=float, default=0.9)
        self.__parser.add_argument('--score', '-s', type=float, default=0.0)
        self.__parser.add_argument('--learning_rate', '-lr', type=float, default=0.0001, dest='l_rate')
        self.__parser.add_argument('--GAMMA', '-G', type=float, default=0.99)

        self.__parser.add_argument('--SEED', '-S', type=int, default=999)
        self.__parser.add_argument('--env_type', '-env', type=str, default='Surviving', choices=['Surviving', 'PP'],
                                   dest='env')
        self.__parser.add_argument('--env_difficulty_level', '-edl', type=str, default='easy',
                                   choices=['easy', 'medium', 'hard', 'harder'])
        self.__parser.add_argument('--model_name', '-m', type=str, default='DGN_R1',
                                   choices=['DGN_R1', 'DGN_R2', 'GACPN_R1', 'GACPN_R2', 'GACPN_R3', 'CommNet',
                                            'TranDGN', 'TranDGN-multi', 'TranDGN-processor', 'TranDGN-tran'])

        self.__surviving_param_dict = {
            'easy': {'map_size': 14, 'comm_dis': 3, 'n_agent': 10, 'max_step': 100, 'n_episode': 3000},
            'medium': {'map_size': 18, 'comm_dis': 4, 'n_agent': 20, 'max_step': 150, 'n_episode': 5000},
            'hard': {'map_size': 24, 'comm_dis': 5, 'n_agent': 30, 'max_step': 200, 'n_episode': 8000},
            'harder': {'map_size': 30, 'comm_dis': 6, 'n_agent': 40, 'max_step': 200, 'n_episode': 6000}
        }

        self.__pp_param_dict = {
            'easy': {'dim': 5, 'vision': 1, 'max_step': 30, 'n_agent': 4, 'n_episode': 5000, 'n_enemies': 1},
            'medium': {'dim': 10, 'vision': 1, 'max_step': 40, 'n_agent': 5, 'n_episode': 8000, 'n_enemies': 1},
            'hard': {'dim': 20, 'vision': 1, 'max_step': 80, 'n_agent': 10, 'n_episode': 12000, 'n_enemies': 1}
        }

        self.__args = None
        self.__env_model_flag = False
        self.__buffer_flag = False

    def execute_args_from_cmd(self, command_argv: List[str]):

        self.__args = self.__parser.parse_args(command_argv)
        if self.__args.env == 'Surviving':
            self.__args.env_param_dict = self.__surviving_param_dict[self.__args.env_difficulty_level]
        elif self.__args.env == 'PP':
            self.__args.env_param_dict = self.__pp_param_dict[self.__args.env_difficulty_level]
        else:
            print(f"Unknown env type{self.__args.env}")
        return self.__args

    def get_env_and_model(self):
        if self.__args is None:
            raise ValueError("args has not been handled, please execute get_args_from_cmd() first!!")

        if self.__env_model_flag:
            raise RuntimeError("env and model has been gained, refuse to create now")
        else:
            self.__env_model_flag = True

        _env = None
        if self.__args.env == 'Surviving':
            _n_agent = self.__args.env_param_dict['n_agent']
            _map_size = self.__args.env_param_dict['map_size']
            _comm_dis = self.__args.env_param_dict['comm_dis']
            _env = Surviving(n_agent=_n_agent, map_size=_map_size, comm_dis=_comm_dis)
        elif self.__args.env == 'PP':
            _dim = self.__args.env_param_dict['dim']
            _vision = self.__args.env_param_dict['vision']
            _max_steps = self.__args.env_param_dict['max_step']
            _n_agent = self.__args.env_param_dict['n_agent']
            _n_enemies = self.__args.env_param_dict['n_enemies']
            _env = PredatorPreyEnv()

            _env.init_curses()
            _env.init_args(self.__parser)

            args = self.__parser.parse_args()
            args.nfriendly = _n_agent
            args.dim = _dim
            args.vision = _vision
            args.nenemies = _n_enemies

            _env.multi_agent_init(args)
        else:
            raise ValueError("Unknown env type {0}".format(self.__args.env))

        _n_agent = _env.n_agent
        _observation_space = _env.len_obs
        _n_actions = _env.n_action
        _hidden_dim = self.__args.hidden_dim
        detached_obs_one_length = _env.detached_obs_one_length * 3

        if self.__args.model_name == 'DGN_R1':
            _model = DGN_R1(_n_agent, _observation_space, _hidden_dim, _n_actions)
            _model_tar = DGN_R1(_n_agent, _observation_space, _hidden_dim, _n_actions)
        elif self.__args.model_name == 'DGN_R2':
            _model = DGN_R2(_n_agent, detached_obs_one_length, _hidden_dim, _n_actions)
            _model_tar = DGN_R2(_n_agent, detached_obs_one_length, _hidden_dim, _n_actions)
        # DGN+MAGIC -> GACPN
        elif self.__args.model_name == 'GACPN_R1':
            _model = GACPN_R1(_n_agent, _observation_space, _hidden_dim, _n_actions)
            _model_tar = GACPN_R1(_n_agent, _observation_space, _hidden_dim, _n_actions)
        elif self.__args.model_name == 'GACPN_R2':
            _model = GACPN_R2(_n_agent, _observation_space, _hidden_dim, _n_actions)
            _model_tar = GACPN_R2(_n_agent, _observation_space, _hidden_dim, _n_actions)
        elif self.__args.model_name == 'GACPN_R3':
            _model = GACPN_R3(_n_agent, detached_obs_one_length, _hidden_dim, _n_actions)
            _model_tar = GACPN_R3(_n_agent, detached_obs_one_length, _hidden_dim, _n_actions)
        # CommNet
        elif self.__args.model_name == 'CommNet':
            _model = CommNet(_n_agent, _observation_space, _hidden_dim, _n_actions)
            _model_tar = CommNet(_n_agent, _observation_space, _hidden_dim, _n_actions)
        # DGN + Transformer -> TranDGN
        elif self.__args.model_name == 'TranDGN':
            _model = TranDGN(_n_agent, _observation_space, _hidden_dim, _n_actions, detached_obs_one_length // 3)
            _model_tar = TranDGN(_n_agent, _observation_space, _hidden_dim, _n_actions, detached_obs_one_length // 3)
        elif self.__args.model_name == 'TranDGN-multi':
            _model = TranDGNWithoutMulti(_n_agent, _observation_space, _hidden_dim, _n_actions,
                                         detached_obs_one_length // 3)
            _model_tar = TranDGNWithoutMulti(_n_agent, _observation_space, _hidden_dim, _n_actions,
                                             detached_obs_one_length // 3)
        elif self.__args.model_name == 'TranDGN-processor':
            _model = TranDGNWithoutProcessor(_n_agent, _observation_space, _hidden_dim, _n_actions,
                                             detached_obs_one_length // 3)
            _model_tar = TranDGNWithoutProcessor(_n_agent, _observation_space, _hidden_dim, _n_actions,
                                                 detached_obs_one_length // 3)
        elif self.__args.model_name == 'TranDGN-tran':
            _model = TranDGNWithoutTran(_n_agent, _observation_space, _hidden_dim, _n_actions,
                                             detached_obs_one_length // 3)
            _model_tar = TranDGNWithoutTran(_n_agent, _observation_space, _hidden_dim, _n_actions,
                                                 detached_obs_one_length // 3)
        else:
            raise TypeError("模型名错误，没有{0}这个模型".format(self.__args.model_name))

        return _env, _model, _model_tar

    def get_buffer(self):
        if self.__args is None:
            raise ValueError("args has not been handled, please execute get_args_from_cmd() first!!")

        if self.__buffer_flag:
            raise RuntimeError('ReplayBuffer has gained, refuse to create new!')
        else:
            self.__buffer_flag = True

        _buffer_size = self.__args.capacity
        _buffer = ReplayBuffer(buffer_size=_buffer_size)
        return _buffer
